import numpy as np
import scipy
import time
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import scipy.ndimage as ndi
import scipy.io
import sys
import scipy.linalg as la

maxiteration = 2000
numdata = 300
D = 3

# state variabales
Rs = np.zeros((maxiteration+1, D, D))
ts = np.zeros((maxiteration+1, D))
logsigma0s = np.zeros(maxiteration+1)
os = np.zeros((maxiteration+1, numdata))
objs = np.zeros(maxiteration+1)
prioroutliers = 4.
priorinliers = 16.

# set parameters
logsigma0mean = -1.
logsigma0var = 0.5
logsigma1mean = 0.

# step parameters
logsigma0step = 0.1
Rstep = 0.05
tstep = 0.01
l = 3
def probcombine(X, Y):
  print X.shape
  print Y.shape
  global Rs, ts, logsigma0s, os, objs
  Rs[0] = np.identity(D)
  ts[0] = np.zeros(D)
  logsigma0 = logsigma0s[0] = 0

  llhs = np.zeros(maxiteration+1)
  llhs[0] = llh(X,Y,np.identity(D),np.zeros(D), np.zeros(numdata), np.exp(-2))
  print 'Init obj = %f log = %f'\
         % (llhs[0], logsigma0)

  rejecto = 0
  for i in range(maxiteration):
    R = np.array(Rs[i]); t = np.array(ts[i]); 
    o = np.array(os[i]); logsigma0 = np.array(logsigma0s[i])
    sigma2 = np.exp(2*logsigma0)
    objs[i] = llh(X,Y,R,t,o,sigma2)
    print 'i = %d obj = %f logsigma = %f sumo = %d rejecto = %d' % (i, objs[i], logsigma0, o.sum(), rejecto)
    
    ###### Update Os ########
    if i % 50 == 0:
     rejecto = 0
     for par in range(numdata):
      old = o[par]
      numout = o.sum()- o[par]
      prioron = (numout + prioroutliers)/(prioroutliers + priorinliers + numdata-1)
      prioroff = 1-prioron
      if old == 1: pp = prioron
      else: pp = prioroff
      orill = llh(X,Y,R,t,o,sigma2) + np.log(pp)
      o[par] = 1 - o[par]
      newll = llh(X,Y,R,t,o,sigma2) + np.log(1.-pp)

      if np.log(np.random.rand()) > newll - orill:
        o[par] = old
        rejecto += 1
    os[i+1] = o

    ####### Update sigma #######
    orilh = llh(X,Y,R,t,o,sigma2) + priorsigma(logsigma0)
    old = logsigma0
    logsigma0 = logsigma0 + logsigma0step * np.random.randn()
    sigma2p = np.exp(2*logsigma0)
    newlh = llh(X,Y,R,t,o,sigma2p) + priorsigma(logsigma0)
    if np.log(np.random.rand()) > newlh - orilh:
      logsigma0 = old
    logsigma0s[i+1] = logsigma0
    sigma2 = np.exp(2*logsigma0)

    ####### Update R #######
    orilh = llh(X,Y,R,t,o,sigma2)
    old = R
    R = R + Rstep * np.random.randn(D,D)
    U,C,V = la.svd(R)
    R = np.dot(U,V)
    newlh = llh(X,Y,R,t,o,sigma2)
    if np.log(np.random.rand()) > newlh - orilh:
      R = old
    Rs[i+1] = R

    ####### Update t #######
    orilh = llh(X,Y,R,t,o,sigma2)
    old = t
    t = t + tstep * np.random.randn(D)
    newlh = llh(X,Y,R,t,o,sigma2)
    if np.log(np.random.rand()) > newlh - orilh:
       t = old
    ts[i+1] = t
  
  maxind = np.argmax(objs)
  return Rs[maxind], ts[maxind][:,None]

def priorsigma(logsigma0):
  return -(logsigma0 - logsigma0mean)*(logsigma0 - logsigma0mean)/ (2.* logsigma0var)


def llh(X,Y,R,t,o,sigma2, moreinfo = False):
  llh = 0
  llhin = 0
  llhout = 0
  Ys = np.dot(R,Y) + t[:,None]
  M = X.shape[1]
  sumo = o.sum()
  outlierprop = (prioroutliers + sumo) / (prioroutliers + M + priorinliers)
  # see pg 439 in Bishop for the likelihood formula
  for c in range(M):
    if o[c] == 0:
      diff = Ys - X[:,c:c+1]
      normal = np.exp(-(diff ** 2).sum(0)/(2*sigma2))/((sigma2*2*np.pi)**(D/2.))
      inside = (normal).mean()*(1.-outlierprop)
      llhin += np.log(inside)
    else:
      llhout += np.log(outlierprop * 1 / (2*l)**D )
  if moreinfo: return llhin, llhout
  return llhin + llhout

